package pl;

import javacard.framework.ISO7816;
import javacard.framework.ISOException;
import javacard.framework.JCSystem;
import javacard.framework.Util;
import javacard.security.MessageDigest;

public class KDF2BytesGenerator {

	private MessageDigest md; 
	
	private byte[] shared;
	private short shLen;
	
	private byte[] iv;
	private short ivLen;
	
	private byte[] C;
	private byte[] hash;

	public KDF2BytesGenerator() {
		
		C = JCSystem.makeTransientByteArray((short) 4,
				JCSystem.CLEAR_ON_DESELECT); 
		hash = JCSystem.makeTransientByteArray((short) 64,
				JCSystem.CLEAR_ON_DESELECT); 
		shared = new byte[(short)544];
	}

	public void init(MessageDigest digest, byte[] shared, short sharLen, byte[] iv, short ivLen) {
		
		Util.arrayCopy(shared, (short)0, this.shared, (short)0, sharLen);
		
		this.shLen = sharLen;
		this.iv = iv;
		this.ivLen = ivLen;
		
		this.md = digest;
		this.md.reset();
	}

	public short generateBytes(byte[] out, short outOff, short len, boolean copyShared) {
		
		if (copyShared) Util.arrayCopy(shared, (short)0, out, outOff, (short) out.length);
		
		if ((short)(out.length - len) < outOff) {
			ISOException.throwIt(ISO7816.SW_WRONG_LENGTH);
		}

		short oBytes = len;
		short outLen = md.getLength();
		short cThreshold = (short) ((short)(oBytes + outLen - 1) / outLen);
		short counter = (short) 1;

		for (short i = 0; i < cThreshold; i++) {
			md.update(shared, (short) 0, shLen);
			
			EMSA_PSS.I2OSP(counter, C);
			if (iv != null) {
				md.update(C, (short) 0, (short) 4);
				md.doFinal(iv, (short) 0, ivLen, hash, (short) 0);
			} else {
				md.doFinal(C, (short) 0, (short) C.length, hash, (short) 0);
				md.reset();
			}

			if (len > outLen) {
				Util.arrayCopy(hash, (short) 0, out, outOff, outLen);
				outOff += outLen;
				len -= outLen;
			} else {
				Util.arrayCopy(hash, (short) 0, out, outOff, len);
			}
			counter++;
		}
		md.reset();
		return len;
	}
}